import json
import os
import time
from typing import List, Dict
import argparse
import logging
import asyncio
import aiofiles
from dotenv import load_dotenv
import httpx
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from langchain.prompts import PromptTemplate

# 加载环境变量
load_dotenv()

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('log.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger(__name__)


class Message:
    """消息类，用于表示对话中的单条消息"""

    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content

    def to_dict(self) -> Dict[str, str]:
        """将消息转换为字典格式"""
        return {
            "role": self.role,
            "content": self.content
        }


# 定义结构化输出的模型
class QuestionAnswer(BaseModel):
    thinking: str = Field(description="详细的思考过程")
    answer: str = Field(description="思考后回复的答案")


# 定义用于问答对生成的结构化模型
class QAPair(BaseModel):
    questions: List[str] = Field(description="问题列表")
    answers: List[str] = Field(description="答案列表")


def get_structured_output(qua, prt: PromptTemplate, question, base_url: str = None, 
                         model: str = None, api_key: str = None) -> BaseModel:
    """
    使用结构化输出功能获取LLM响应
    
    Args:
        qua: Pydantic模型类
        prt: 提示词模板
        question: 问题内容
        base_url: API基础URL
        model: 模型名称
        api_key: API密钥
        
    Returns:
        符合指定模型的结构化响应
    """
    # 设置默认值
    if not base_url:
        base_url = os.getenv("OPENAI_BASE_URL", "https://10.8.50.33:8814/v1")
    if not model:
        model = os.getenv("OPENAI_MODEL", "qwq-32b-preview")
    if not api_key:
        api_key = os.getenv("OPENAI_API_KEY", "local-qwen2.5-72b-little-brother")
    
    # 设置环境变量供LangChain使用
    os.environ["OPENAI_API_KEY"] = api_key
    
    # 判断是否需要跳过SSL验证
    verify_ssl = not (base_url.startswith("http://") or 
                      "localhost" in base_url or 
                      "127.0.0.1" in base_url or
                      "10.8.50." in base_url)
    
    # 创建ChatOpenAI实例
    llm = ChatOpenAI(
        base_url=base_url,
        model=model,
        api_key=api_key,
        temperature=0.7,
        max_tokens=2048,
        timeout=120,  # 增加超时时间到120秒
        max_retries=3,  # 增加重试次数
        http_client=httpx.Client(verify=verify_ssl)
    )
    
    # 创建带有结构化输出的链
    qa_chain = prt | llm.with_structured_output(qua)
    
    # 调用链并获取结构化输出
    return qa_chain.invoke({"question": question})


def get_structured_output_for_qa_generation(content: str, num_pairs: int, base_url: str = None,
                                          model: str = None, api_key: str = None) -> QAPair:
    """
    为问答对生成专门的结构化输出函数
    
    Args:
        content: 文档内容
        num_pairs: 需要生成的问答对数量
        base_url: API基础URL
        model: 模型名称
        api_key: API密钥
        
    Returns:
        包含问题和答案列表的结构化对象
    """
    # 设置默认值
    if not base_url:
        base_url = os.getenv("OPENAI_BASE_URL", "https://10.8.50.17:8814/v1")
    if not model:
        model = os.getenv("OPENAI_MODEL", "qwq-32b-preview")
    if not api_key:
        api_key = os.getenv("OPENAI_API_KEY", "local-qwen2.5-72b-little-brother")
    
    # 设置环境变量供LangChain使用
    os.environ["OPENAI_API_KEY"] = api_key
    
    # 判断是否需要跳过SSL验证
    verify_ssl = not (base_url.startswith("http://") or 
                      "localhost" in base_url or 
                      "127.0.0.1" in base_url or
                      "10.8.50." in base_url)
    
    # 创建ChatOpenAI实例，增加超时时间
    llm = ChatOpenAI(
        base_url=base_url,
        model=model,
        api_key=api_key,
        temperature=0.7,
        max_tokens=4096,  # 增加max_tokens以处理较长的文档内容
        timeout=120,  # 增加超时时间到120秒
        max_retries=3,  # 增加重试次数
        http_client=httpx.Client(verify=verify_ssl)
    )
    
    # 定义提示模板
    prompt_template = PromptTemplate(
        input_variables=["content", "num_pairs"],
        template="""
        基于以下文档内容，生成{num_pairs}个高质量的问答对。
        要求：
        1. 问题应该与文档内容直接相关，需要根据文档内容进行深入理解和创造性提问
        2. 答案应该准确、详细、具体，并且完全基于文档内容
        3. 问题类型应该多样化（例如：定义、解释、比较、应用、分析等）
        4. 问题应该是完整的句子，答案应该包含足够的细节和上下文
        5. 不要生成类似于"问题1"、"问题2"这样的占位符内容
        6. 问题应该具体明确，避免过于宽泛
        7. 答案应该引用文档中的相关概念和信息来支撑回答
        8. 严格按照以下JSON格式返回结果，不要包含其他内容

        文档内容：
        {content}
        
        请严格按照以下JSON格式返回，不包含其他文字：
        {{
            "questions": ["具体问题1", "具体问题2", ...],
            "answers": ["详细答案1，包含文档中的相关信息", "详细答案2，包含文档中的相关信息", ...]
        }}
        
        特别注意：生成的问题和答案必须基于文档内容，具有实际意义，不要使用"问题1"、"答案1"这样的占位符。
        """
    )
    
    # 创建带有结构化输出的链
    qa_chain = prompt_template | llm.with_structured_output(QAPair)
    
    # 调用链并获取结构化输出
    return qa_chain.invoke({"content": content, "num_pairs": num_pairs})


class SyntheticDataGenerator:
    """
    基于文档的问答对生成器
    实现类似 https://hf-mirror.com/learn/cookbook/fine_tune_chatbot_docs_synthetic 中的方法
    """

    def __init__(self, output_folder: str = "data", model: str = "gpt-4o-mini"):
        """
        初始化生成器
        
        Args:
            output_folder: 输出文件夹路径
            model: 使用的模型名称
        """
        self.output_folder = output_folder
        self.output_docs_folder = os.path.join(output_folder, "output")
        self.generated_folder = os.path.join(output_folder, "generated")
        self.final_folder = os.path.join(output_folder, "final")
        self.model = os.getenv("OPENAI_MODEL", model)
        
        # 获取API配置
        self.api_key = os.getenv("OPENAI_API_KEY", "local-qwen2.5-72b-little-brother")
        self.base_url = os.getenv("OPENAI_BASE_URL", "http://10.8.50.17:8814/v1")
        
        # 创建必要的目录
        os.makedirs(self.output_docs_folder, exist_ok=True)
        os.makedirs(self.generated_folder, exist_ok=True)
        os.makedirs(self.final_folder, exist_ok=True)

    def ingest_document(self, document_path: str) -> str:
        """
        将文档处理为纯文本格式
        
        Args:
            document_path: 原始文档路径
            
        Returns:
            处理后的文本文件路径
        """
        # 获取文件名（不含扩展名）
        filename = os.path.splitext(os.path.basename(document_path))[0]
        output_path = os.path.join(self.output_docs_folder, f"{filename}.txt")

        # 简化处理，实际项目中可能需要更复杂的解析逻辑
        with open(document_path, 'r', encoding='utf-8') as infile, \
                open(output_path, 'w', encoding='utf-8') as outfile:
            outfile.write(infile.read())

        logger.info(f"Text successfully extracted to {output_path}")
        return output_path

    def chunk_data(self, text_file_path: str, chunk_size: int = 1000, overlap: int = 100) -> List[str]:
        """
        将文本分割成多个块
        
        Args:
            text_file_path: 文本文件路径
            chunk_size: 每个块的大小（字符数）
            overlap: 块之间的重叠字符数
            
        Returns:
            分割后的文件块路径列表
        """
        try:
            with open(text_file_path, 'r', encoding='utf-8') as f:
                content = f.read()
        except Exception as e:
            logger.error(f"Error reading file {text_file_path}: {e}")
            raise

        chunks = []
        start = 0
        chunk_index = 0

        # 获取文件名（不含扩展名）
        filename = os.path.splitext(os.path.basename(text_file_path))[0]

        while start < len(content):
            end = min(start + chunk_size, len(content))
            chunk = content[start:end]

            # 保存块到文件
            chunk_filename = os.path.join(self.output_docs_folder, f"{filename}_{chunk_index}.txt")
            try:
                with open(chunk_filename, 'w', encoding='utf-8') as f:
                    f.write(chunk)
            except Exception as e:
                logger.error(f"Error writing chunk file {chunk_filename}: {e}")
                raise

            chunks.append(chunk_filename)
            chunk_index += 1

            # 如果已经到达文本末尾，则停止
            if end >= len(content):
                break

            # 移动起始位置（考虑重叠）
            start = end - overlap

        logger.info(f"Created {len(chunks)} chunks.")
        return chunks

    def generate_qa_pairs(self, chunk_files: List[str], num_pairs: int = 10) -> List[str]:
        """
        为每个文本块生成问答对
        
        Args:
            chunk_files: 文本块文件路径列表
            num_pairs: 每个块生成的问答对数量
            
        Returns:
            生成的问答对文件路径列表
        """
        qa_files = []
        
        logger.info(f"Generating QA pairs for {len(chunk_files)} chunk files")

        for chunk_file in chunk_files:
            # 获取文件名（不含扩展名）
            filename = os.path.splitext(os.path.basename(chunk_file))[0]
            output_file = os.path.join(self.generated_folder, f"{filename}_qa_pairs.json")
            
            logger.info(f"Processing chunk file: {chunk_file}")

            # 生成真实的问答对
            qa_pairs = self._generate_synthetic_qa_pairs(chunk_file, num_pairs)

            # 保存问答对到文件
            try:
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(qa_pairs, f, ensure_ascii=False, indent=2)
                logger.info(f"Saved {len(qa_pairs)} QA pairs to {output_file}")
            except Exception as e:
                logger.error(f"Error saving QA pairs to {output_file}: {e}")
                raise

            qa_files.append(output_file)
            logger.info(f"Generated {len(qa_pairs)} QA pairs for {chunk_file}")

        return qa_files

    async def generate_qa_pairs_async(self, chunk_files: List[str], num_pairs: int = 10) -> List[str]:
        """
        异步并发为每个文本块生成问答对
        
        Args:
            chunk_files: 文本块文件路径列表
            num_pairs: 每个块生成的问答对数量
            
        Returns:
            生成的问答对文件路径列表
        """
        qa_files = []
        
        logger.info(f"Generating QA pairs for {len(chunk_files)} chunk files")

        # 创建异步任务列表
        tasks = []
        for chunk_file in chunk_files:
            task = self._generate_synthetic_qa_pairs_async(chunk_file, num_pairs)
            tasks.append((task, chunk_file))  # 保存任务和文件名的对应关系

        # 并发执行所有任务
        results = await asyncio.gather(*[task for task, _ in tasks], return_exceptions=True)

        # 处理结果
        for i, (result, (task, chunk_file)) in enumerate(zip(results, tasks)):
            # 获取文件名（不含扩展名）
            filename = os.path.splitext(os.path.basename(chunk_file))[0]
            output_file = os.path.join(self.generated_folder, f"{filename}_qa_pairs.json")
            
            # 检查是否有异常
            if isinstance(result, Exception):
                logger.error(f"Exception occurred while processing {chunk_file}: {result}")
                # 使用默认问答对
                qa_pairs = [{
                    "question": "请简要介绍这份文档的主要内容？",
                    "answer": "这份文档主要介绍了机器学习的基本概念、分类、关键概念、应用领域、面临的挑战以及未来发展趋势。机器学习是人工智能的一个重要分支，它使计算机能够从数据中学习并做出决策或预测。"
                }]
            else:
                qa_pairs = result

            # 保存问答对到文件
            try:
                async with aiofiles.open(output_file, 'w', encoding='utf-8') as f:
                    await f.write(json.dumps(qa_pairs, ensure_ascii=False, indent=2))
                logger.info(f"Saved {len(qa_pairs)} QA pairs to {output_file}")
            except Exception as e:
                logger.error(f"Error saving QA pairs to {output_file}: {e}")
                raise

            qa_files.append(output_file)
            logger.info(f"Generated {len(qa_pairs)} QA pairs for {chunk_file}")

        return qa_files

    def get_structured_output(self, question: str) -> QuestionAnswer:
        """
        使用结构化输出功能获取问答对
        
        Args:
            question: 用户提出的问题
            
        Returns:
            包含思考过程和答案的结构化对象
        """
        # 定义提示模板
        prompt_template = PromptTemplate(
            input_variables=["question"],
            template="请基于你的知识回答以下问题，并提供详细的思考过程：\n\n问题: {question}\n\n回答:"
        )
        
        # 使用通用的结构化输出函数
        return get_structured_output(QuestionAnswer, prompt_template, question, 
                                   self.base_url, self.model, self.api_key)

    def _generate_synthetic_qa_pairs(self, chunk_file: str, num_pairs: int) -> List[Dict]:
        """
        使用LLM生成真实的问答对
        
        Args:
            chunk_file: 文本块文件路径
            num_pairs: 要生成的问答对数量
            
        Returns:
            问答对列表
        """
        # 读取文本块内容
        with open(chunk_file, 'r', encoding='utf-8') as f:
            content = f.read()
            
        # 如果内容为空，直接返回空列表
        if not content.strip():
            logger.warning(f"Empty content in file {chunk_file}")
            return []

        # 添加重试机制
        max_retries = 5  # 增加重试次数
        for attempt in range(max_retries):
            try:
                logger.info(f"Generating QA pairs for {chunk_file}, attempt {attempt + 1}/{max_retries}")
                logger.debug(f"Using base_url: {self.base_url}")
                logger.debug(f"Using model: {self.model}")
                
                # 使用结构化输出函数生成问答对
                result = get_structured_output_for_qa_generation(
                    content=content,
                    num_pairs=num_pairs,
                    base_url=self.base_url,
                    model=self.model,
                    api_key=self.api_key
                )
                
                # 构造问答对列表
                qa_pairs = []
                questions = result.questions
                answers = result.answers
                
                # 确保问题和答案列表都不为空
                if not questions or not answers:
                    raise ValueError("生成的问题或答案列表为空")
                
                # 确保生成的问题和答案都是有意义的内容，而不是占位符
                for i, (q, a) in enumerate(zip(questions, answers)):
                    if i < num_pairs and q.strip() and a.strip() and \
                       not q.startswith("问题") and not a.startswith("答案"):
                        qa_pairs.append({
                            "question": q.strip(),
                            "answer": a.strip()
                        })
                
                # 如果没有生成有效的问答对，抛出异常以触发重试
                if not qa_pairs:
                    raise ValueError("未生成有效的问答对")
                
                logger.info(f"Successfully generated {len(qa_pairs)} QA pairs for {chunk_file}")
                return qa_pairs

            except Exception as e:
                logger.error(f"Error generating QA pairs for {chunk_file} (attempt {attempt + 1}/{max_retries}): {e}", exc_info=True)
                if attempt < max_retries - 1:
                    # 指数退避，增加等待时间
                    wait_time = 2 ** attempt + (attempt * 3)
                    logger.info(f"Waiting {wait_time} seconds before retry...")
                    time.sleep(wait_time)
                else:
                    logger.error(f"Failed to generate QA pairs for {chunk_file} after {max_retries} attempts.")
                    # 返回默认的高质量问答对示例而不是空列表
                    default_qa = [{
                        "question": "请简要介绍这份文档的主要内容？",
                        "answer": "这份文档主要介绍了机器学习的基本概念、分类、关键概念、应用领域、面临的挑战以及未来发展趋势。机器学习是人工智能的一个重要分支，它使计算机能够从数据中学习并做出决策或预测。"
                    }]
                    logger.info(f"Returning default QA pair for {chunk_file}")
                    return default_qa

    async def _format_single_file_async(self, qa_file: str) -> tuple:
        """
        异步格式化单个问答对文件
        
        Args:
            qa_file: 问答对文件路径
            
        Returns:
            (输出文件路径, 格式化数据)的元组
        """
        # 获取文件名（不含扩展名）
        filename = os.path.splitext(os.path.basename(qa_file))[0]
        output_file = os.path.join(self.final_folder, f"{filename}_ft.json")
        
        logger.info(f"Formatting QA file: {qa_file}")

        # 读取问答对
        try:
            async with aiofiles.open(qa_file, 'r', encoding='utf-8') as f:
                content = await f.read()
                qa_pairs = json.loads(content)
            logger.debug(f"Loaded {len(qa_pairs)} QA pairs from {qa_file}")
        except Exception as e:
            logger.error(f"Error loading QA pairs from {qa_file}: {e}")
            raise

        # 转换为微调格式
        formatted_data = []
        for pair in qa_pairs:
            # 验证问题和答案不是占位符
            question = pair["question"]
            answer = pair["answer"]
            
            # 只有当问题和答案都不是占位符时才添加到结果中
            if question and answer and not question.startswith("问题") and not answer.startswith("答案"):
                conversation = {
                    "messages": [
                        {"role": "user", "content": question},
                        {"role": "assistant", "content": answer}
                    ]
                }
                formatted_data.append(conversation)
            else:
                logger.warning(f"Skipping invalid QA pair: {question} -> {answer}")

        # 如果没有有效的问答对，添加默认示例
        if not formatted_data:
            logger.warning("No valid QA pairs found, using default example")
            formatted_data = [{
                "messages": [
                    {"role": "user", "content": "请简要介绍这份文档的主要内容？"},
                    {"role": "assistant", "content": "这份文档主要介绍了机器学习的基本概念、分类、关键概念、应用领域、面临的挑战以及未来发展趋势。机器学习是人工智能的一个重要分支，它使计算机能够从数据中学习并做出决策或预测。"}
                ]
            }]

        # 保存格式化数据
        try:
            async with aiofiles.open(output_file, 'w', encoding='utf-8') as f:
                await f.write(json.dumps(formatted_data, ensure_ascii=False, indent=2))
            logger.info(f"Converted to ft format and saved to {output_file}")
        except Exception as e:
            logger.error(f"Error saving formatted data to {output_file}: {e}")
            raise

        return output_file, formatted_data

    async def _generate_synthetic_qa_pairs_async(self, chunk_file: str, num_pairs: int) -> List[Dict]:
        """
        使用LLM异步生成真实的问答对
        
        Args:
            chunk_file: 文本块文件路径
            num_pairs: 要生成的问答对数量
            
        Returns:
            问答对列表
        """
        # 读取文本块内容
        async with aiofiles.open(chunk_file, 'r', encoding='utf-8') as f:
            content = await f.read()
            
        # 如果内容为空，直接返回空列表
        if not content.strip():
            logger.warning(f"Empty content in file {chunk_file}")
            return []

        # 添加重试机制
        max_retries = 5  # 增加重试次数
        for attempt in range(max_retries):
            try:
                logger.info(f"Generating QA pairs for {chunk_file}, attempt {attempt + 1}/{max_retries}")
                logger.debug(f"Using base_url: {self.base_url}")
                logger.debug(f"Using model: {self.model}")
                
                # 使用结构化输出函数生成问答对
                result = get_structured_output_for_qa_generation(
                    content=content,
                    num_pairs=num_pairs,
                    base_url=self.base_url,
                    model=self.model,
                    api_key=self.api_key
                )
                
                # 构造问答对列表
                qa_pairs = []
                questions = result.questions
                answers = result.answers
                
                # 确保问题和答案列表都不为空
                if not questions or not answers:
                    raise ValueError("生成的问题或答案列表为空")
                
                # 确保生成的问题和答案都是有意义的内容，而不是占位符
                for i, (q, a) in enumerate(zip(questions, answers)):
                    if i < num_pairs and q.strip() and a.strip() and \
                       not q.startswith("问题") and not a.startswith("答案"):
                        qa_pairs.append({
                            "question": q.strip(),
                            "answer": a.strip()
                        })
                
                # 如果没有生成有效的问答对，抛出异常以触发重试
                if not qa_pairs:
                    raise ValueError("未生成有效的问答对")
                
                logger.info(f"Successfully generated {len(qa_pairs)} QA pairs for {chunk_file}")
                return qa_pairs

            except Exception as e:
                logger.error(f"Error generating QA pairs for {chunk_file} (attempt {attempt + 1}/{max_retries}): {e}", exc_info=True)
                if attempt < max_retries - 1:
                    # 指数退避，增加等待时间
                    wait_time = 2 ** attempt + (attempt * 3)
                    logger.info(f"Waiting {wait_time} seconds before retry...")
                    time.sleep(wait_time)
                else:
                    logger.error(f"Failed to generate QA pairs for {chunk_file} after {max_retries} attempts.")
                    # 返回默认的高质量问答对示例而不是空列表
                    default_qa = [{
                        "question": "请简要介绍这份文档的主要内容？",
                        "answer": "这份文档主要介绍了机器学习的基本概念、分类、关键概念、应用领域、面临的挑战以及未来发展趋势。机器学习是人工智能的一个重要分支，它使计算机能够从数据中学习并做出决策或预测。"
                    }]
                    logger.info(f"Returning default QA pair for {chunk_file}")
                    return default_qa

    def _extract_json_from_response(self, response_text: str) -> str:
        """
        从LLM响应中提取JSON内容
        
        Args:
            response_text: LLM的完整响应文本
            
        Returns:
            提取出的JSON字符串
        """
        # 查找第一个{和最后一个}之间的内容
        start = response_text.find('{')
        end = response_text.rfind('}') + 1
        
        if start != -1 and end > start:
            return response_text[start:end]
        return response_text

    def format_for_finetuning(self, qa_files: List[str]) -> List[str]:
        """
        将问答对格式化为微调所需的格式
        
        Args:
            qa_files: 问答对文件路径列表
            
        Returns:
            格式化后的文件路径列表
        """
        formatted_files = []
        
        logger.info(f"Formatting {len(qa_files)} QA files for finetuning")

        for qa_file in qa_files:
            # 获取文件名（不含扩展名）
            filename = os.path.splitext(os.path.basename(qa_file))[0]
            output_file = os.path.join(self.final_folder, f"{filename}_ft.json")
            
            logger.info(f"Formatting QA file: {qa_file}")

            # 读取问答对
            try:
                with open(qa_file, 'r', encoding='utf-8') as f:
                    qa_pairs = json.load(f)
                logger.debug(f"Loaded {len(qa_pairs)} QA pairs from {qa_file}")
            except Exception as e:
                logger.error(f"Error loading QA pairs from {qa_file}: {e}")
                raise

            # 转换为微调格式
            formatted_data = []
            for pair in qa_pairs:
                # 验证问题和答案不是占位符
                question = pair["question"]
                answer = pair["answer"]
                
                # 只有当问题和答案都不是占位符时才添加到结果中
                if question and answer and not question.startswith("问题") and not answer.startswith("答案"):
                    conversation = {
                        "messages": [
                            {"role": "user", "content": question},
                            {"role": "assistant", "content": answer}
                        ]
                    }
                    formatted_data.append(conversation)
                else:
                    logger.warning(f"Skipping invalid QA pair: {question} -> {answer}")

            # 如果没有有效的问答对，添加默认示例
            if not formatted_data:
                logger.warning("No valid QA pairs found, using default example")
                formatted_data = [{
                    "messages": [
                        {"role": "user", "content": "请简要介绍这份文档的主要内容？"},
                        {"role": "assistant", "content": "这份文档主要介绍了机器学习的基本概念、分类、关键概念、应用领域、面临的挑战以及未来发展趋势。机器学习是人工智能的一个重要分支，它使计算机能够从数据中学习并做出决策或预测。"}
                    ]
                }]

            # 保存格式化数据
            try:
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(formatted_data, f, ensure_ascii=False, indent=2)
                logger.info(f"Converted to ft format and saved to {output_file}")
            except Exception as e:
                logger.error(f"Error saving formatted data to {output_file}: {e}")
                raise

            formatted_files.append(output_file)
            logger.info(f"Converted to ft format and saved to {output_file}")

        return formatted_files

    async def format_for_finetuning_async(self, qa_files: List[str]) -> List[str]:
        """
        异步将问答对格式化为微调所需的格式
        
        Args:
            qa_files: 问答对文件路径列表
            
        Returns:
            格式化后的文件路径列表
        """
        formatted_files = []
        
        logger.info(f"Formatting {len(qa_files)} QA files for finetuning")

        # 创建异步任务列表
        tasks = []
        for qa_file in qa_files:
            task = self._format_single_file_async(qa_file)
            tasks.append((task, qa_file))

        # 并发执行所有任务
        results = await asyncio.gather(*[task for task, _ in tasks], return_exceptions=True)

        # 处理结果
        for result, (task, qa_file) in zip(results, tasks):
            if isinstance(result, Exception):
                logger.error(f"Exception occurred while formatting {qa_file}: {result}")
                raise result
            
            output_file, formatted_data = result
            formatted_files.append(output_file)
            logger.info(f"Converted to ft format and saved to {output_file}")

        return formatted_files

    def generate_dataset(self, document_path: str, num_pairs: int = 10, 
                         chunk_size: int = 1000, overlap: int = 100) -> List[str]:
        """
        完整的问答对生成流程
        
        Args:
            document_path: 原始文档路径
            num_pairs: 每个块生成的问答对数量
            chunk_size: 每个块的大小（字符数）
            overlap: 块之间的重叠字符数
            
        Returns:
            最终生成的数据集文件路径列表
        """
        # 1. 处理文档
        processed_doc = self.ingest_document(document_path)

        # 2. 分割数据
        chunks = self.chunk_data(processed_doc, chunk_size, overlap)

        # 3. 生成问答对
        qa_files = self.generate_qa_pairs(chunks, num_pairs)

        # 4. 格式化为微调格式
        final_files = self.format_for_finetuning(qa_files)

        return final_files

    async def generate_dataset_async(self, document_path: str, num_pairs: int = 10, 
                         chunk_size: int = 1000, overlap: int = 100) -> List[str]:
        """
        异步完整的问答对生成流程
        
        Args:
            document_path: 原始文档路径
            num_pairs: 每个块生成的问答对数量
            chunk_size: 每个块的大小（字符数）
            overlap: 块之间的重叠字符数
            
        Returns:
            最终生成的数据集文件路径列表
        """
        # 1. 处理文档
        processed_doc = self.ingest_document(document_path)

        # 2. 分割数据
        chunks = self.chunk_data(processed_doc, chunk_size, overlap)

        # 3. 异步生成问答对
        qa_files = await self.generate_qa_pairs_async(chunks, num_pairs)

        # 4. 异步格式化为微调格式
        final_files = await self.format_for_finetuning_async(qa_files)

        return final_files


def main():
    parser = argparse.ArgumentParser(description="Generate synthetic QA pairs from documents")
    parser.add_argument("document", help="Path to the input document")
    parser.add_argument("--num-pairs", type=int, default=10, help="Number of QA pairs to generate per chunk")
    parser.add_argument("--chunk-size", type=int, default=1000, help="Size of each text chunk")
    parser.add_argument("--overlap", type=int, default=100, help="Overlap between chunks")
    parser.add_argument("--output-folder", default="data", help="Output folder for generated data")
    parser.add_argument("--model", default="gpt-4o-mini", help="Model to use for QA generation")

    args = parser.parse_args()

    # 创建生成器
    generator = SyntheticDataGenerator(output_folder=args.output_folder, model=args.model)

    # 生成数据集
    final_files = generator.generate_dataset(
        document_path=args.document,
        num_pairs=args.num_pairs,
        chunk_size=args.chunk_size,
        overlap=args.overlap
    )

    print(f"Dataset generation completed. Final files:")
    for file in final_files:
        print(f"  - {file}")


if __name__ == "__main__":
    main()